import cv2
from queue import PriorityQueue
import pickle
import math
import copy
import numpy as np
from PIL import Image
from pyheatmap.heatmap import HeatMap
import random


# 读取地图数据
with open('road_signs_loc_data.pkl', 'rb') as f:  # 读取road_signs_loc
    road_signs_loc = pickle.load(f)
with open('road_signs_neighbor_data.pkl', 'rb') as f:  # 读取road_signs_neighbor
    road_signs_neighbor = pickle.load(f)
with open('people_loc_data.pkl', 'rb') as f:  # 读取people_loc
    people_loc = pickle.load(f)
with open('obs_people_loc_data.pkl', 'rb') as f:  # 读取obs_people_loc
    obs_people_loc = pickle.load(f)
with open('safety_loc_data.pkl', 'rb') as f:  # 读取safety_loc
    safety_loc = pickle.load(f)
img_square_org = cv2.imread('img_square_color.png')  # 彩色地图
img_square_black_and_white = cv2.imread('img_square_black_and_white.png')  # 黑白地图通道
img_square_signs = cv2.imread('img_square_signs.png')  # 路标地图
img_square_safety = cv2.imread('img_square_safety.png')  # 安全出口地图


class Graph:
    @staticmethod
    def neighbors(sign):
        return road_signs_neighbor[str(sign)]['neighbour']

    @staticmethod
    def cost(current, ind):
        return road_signs_neighbor[str(current)]['time'][ind]

    @staticmethod
    def heuristic(s1, s2):
        return abs(s2[0]-s1[0]) + abs(s2[1]-s1[1])  # |x1-x2| + |y1-y2|


# A_star算法
def a_star(start, goal):
    frontier = PriorityQueue()
    frontier.put(start, 0)
    came_from = dict()
    cost_so_far = dict()
    came_from[str(start)] = None
    cost_so_far[str(start)] = 0

    while not frontier.empty():
        current = frontier.get()
        if current == goal:
            break

        for ind, next in enumerate(Graph.neighbors(current)):
            new_cost = cost_so_far[str(current)] + Graph.cost(current, ind)  # Graph.cost(current, next)
            if str(next) not in cost_so_far or new_cost < cost_so_far[str(next)]:
                cost_so_far[str(next)] = new_cost
                priority = new_cost + Graph.heuristic(goal, next)
                frontier.put(next, priority)
                came_from[str(next)] = current

    return came_from


# 路径绘制函数
def draw_path(map, path):
    for i in range(0, len(path)-1):
        s1, s2 = tuple(path[i]), tuple(path[i+1])
        cv2.line(map, s1, s2, (255, 255, 0), 7)  # （翠蓝）


# 逆寻路径函数
def path_find(start, goal):
    came_from = a_star(start, goal)
    path, current = [], goal
    while current != start:
        path.append(current)
        current = came_from[str(current)]
    path.append(start)
    path.reverse()
    return path


# 欧式距离
def distance_enclid(s1, s2):
    return math.sqrt(sum([(a - b)**2 for (a, b) in zip(s1, s2)]))


# 初始化路标间通行时间
def initial_path_time(road_signs_neighbor):
    for _, value in road_signs_neighbor.items():
        for i in range(len(value['cost'])):
            value['time'].append(round((value['cost'][i] / 1.7), 3))  # 保留三位小数
    return road_signs_neighbor


# 随机生成障碍人群集合
def random_obs_people(num, road_signs_loc, people_loc):
    obs_people = []
    count = 0
    while True:
        if count == num:
            return obs_people
        ind = np.random.randint(0, len(road_signs_loc))
        if (road_signs_loc[ind] not in people_loc) and (road_signs_loc[ind] not in obs_people):
            obs_people.append(road_signs_loc[ind])
            count += 1


# 根据当前障碍人群集合位置，更新路标间通行时间
def update_road_signs_neighbor(obs_people_loc, road_signs_neighbor_org):
    for obs_loc in obs_people_loc:
        for i, neighbor in enumerate(road_signs_neighbor_org[str(obs_loc)]['neighbour']):  # 遍历该位置路标旁的所有neighbor
            if neighbor in obs_people_loc:  # 若存在邻居sign也被obs_people集合占据，更改两点间路径通过时间
                road_signs_neighbor_org[str(obs_loc)]['time'][i] *= 7  # 路径通过时间 *7
    return road_signs_neighbor_org


def apply_heatmap(image, data):
    '''image是原图，data是坐标'''
    '''创建一个新的与原图大小一致的图像，color为0背景为黑色。这里这样做是因为在绘制热力图的时候如果不选择背景图，画出来的图与原图大小不一致（根据点的坐标来的），导致无法对热力图和原图进行加权叠加，因此，这里我新建了一张背景图。'''
    background = Image.new("RGB", (image.shape[1], image.shape[0]), color=0)
    # 开始绘制热度图
    hm = HeatMap(data)
    hit_img = hm.heatmap(base=background, r=70)  # background为背景图片，r是半径，默认为10
    hit_img = cv2.cvtColor(np.asarray(hit_img), cv2.COLOR_RGB2BGR)  # Image格式转换成cv2格式
    overlay = image.copy()
    alpha = 0.5  # 设置覆盖图片的透明度
    # cv2.rectangle(overlay, (0, 0), (image.shape[1], image.shape[0]), (255, 0, 0), -1)  # 设置蓝色为热度图基本色蓝色
    # image = cv2.addWeighted(overlay, alpha, image, 1-alpha, 0)  # 将背景热度图覆盖到原图
    image = cv2.addWeighted(hit_img, alpha, image, 1-alpha, 0)  # 将热度图覆盖到原图
    return image


draw_obs_loc = []


def cross_growth(point, dis, image):
    x, y = point[0], point[1]
    p1, p2, p3, p4 = [x+dis, y], [x, y+dis], [x-dis, y], [x, y-dis]
    cv2.circle(image, (p1[0], p1[1]), 1, (0, 0, 255), -1)
    cv2.circle(image, (p2[0], p2[1]), 1, (0, 0, 255), -1)
    cv2.circle(image, (p3[0], p3[1]), 1, (0, 0, 255), -1)
    cv2.circle(image, (p4[0], p4[1]), 1, (0, 0, 255), -1)


def x_growth(point, dis, image):
    x, y = point[0], point[1]
    p1, p2, p3, p4 = [x+dis, y+dis], [x-dis, y+dis], [x-dis, y-dis], [x+dis, y-dis]
    cv2.circle(image, (p1[0], p1[1]), 1, (0, 0, 255), -1)
    cv2.circle(image, (p2[0], p2[1]), 1, (0, 0, 255), -1)
    cv2.circle(image, (p3[0], p3[1]), 1, (0, 0, 255), -1)
    cv2.circle(image, (p4[0], p4[1]), 1, (0, 0, 255), -1)


def draw_point_increase(image):
    for obs_loc in obs_people_loc:
        neighbors = road_signs_neighbor[str(obs_loc)]['neighbour']
        for n_loc in neighbors:  # 找到obs_loc邻居路标中，也处于障碍合集obs_people_loc中的路标
            if n_loc in obs_people_loc:
                cv2.circle(image, (obs_loc[0], obs_loc[1]), 1, (0, 0, 255), -1)
                cv2.circle(image, (n_loc[0], n_loc[1]), 1, (0, 0, 255), -1)

                x1, x2, y1, y2 = obs_loc[0], n_loc[0], obs_loc[1], n_loc[1]

                mid_2_loc = [(x1 + x2) // 2, (y1 + y2) // 2]
                mid_1_loc = [(x1 + mid_2_loc[0]) // 2, (y1 + mid_2_loc[1]) // 2]
                mid_3_loc = [(x2 + mid_2_loc[0]) // 2, (y2 + mid_2_loc[1]) // 2]

                cv2.circle(image, (mid_1_loc[0], mid_1_loc[1]), 1, (0, 0, 255), -1)
                cv2.circle(image, (mid_2_loc[0], mid_2_loc[1]), 1, (0, 0, 255), -1)
                cv2.circle(image, (mid_3_loc[0], mid_3_loc[1]), 1, (0, 0, 255), -1)

                # cross_growth(mid_1_loc, 5, image)
                # cross_growth(mid_2_loc, 5, image)
                # cross_growth(mid_3_loc, 5, image)

                x_growth(mid_1_loc, 5, image)
                x_growth(mid_2_loc, 5, image)
                x_growth(mid_3_loc, 5, image)
            # else:
            #     neighbors_neighbors = road_signs_neighbor[str(n_loc)]['neighbour']
            #     for n_n_loc in neighbors_neighbors:
            #         if n_n_loc in obs_people_loc:
            #             cv2.circle(image, (n_n_loc[0], n_n_loc[1]), 1, (0, 0, 255), -1)
            #
            #             x1, x2, y1, y2 = obs_loc[0], n_n_loc[0], obs_loc[1], n_n_loc[1]
            #
            #             mid_2_loc = [(x1 + x2) // 2, (y1 + y2) // 2]
            #             mid_1_loc = [(x1 + mid_2_loc[0]) // 2, (y1 + mid_2_loc[1]) // 2]
            #             mid_3_loc = [(x2 + mid_2_loc[0]) // 2, (y2 + mid_2_loc[1]) // 2]
            #
            #             cv2.circle(image, (mid_1_loc[0], mid_1_loc[1]), 1, (0, 0, 255), -1)
            #             cv2.circle(image, (mid_2_loc[0], mid_2_loc[1]), 1, (0, 0, 255), -1)
            #             cv2.circle(image, (mid_3_loc[0], mid_3_loc[1]), 1, (0, 0, 255), -1)
            #
            #             cross_growth(mid_1_loc, 3, image)
            #             cross_growth(mid_2_loc, 3, image)
            #             cross_growth(mid_3_loc, 3, image)
            #
            #             x_growth(mid_1_loc, 3, image)
            #             x_growth(mid_2_loc, 3, image)
            #             x_growth(mid_3_loc, 3, image)


"""
# 设置起点、终点
start, goal = road_signs_loc[0], road_signs_loc[-1]

# 得到最短路径
path = path_find(start, goal)

# 画出最短路径
draw_path(img_square_org, path)
"""

def main():
    global people_loc, road_signs_neighbor, obs_people_loc
    people_loc_new = []  # 用于更新客户人群的新位置
    obs_people_loc_new = []  # 用于更新障碍人群的新位置

    reach_count = 0  # 记录已经到达safety的客户人群数量
    people_group_count = len(people_loc)  # 客户人群数量

    # 初始化路标间通行时间（仅参考距离因素）
    road_signs_neighbor = initial_path_time(road_signs_neighbor)
    road_signs_neighbor_org = copy.deepcopy(road_signs_neighbor)  # 保存初始road_signs_neighbor信息

    # 初始化障碍人群集合
    obs_people_loc = random_obs_people(120, road_signs_loc, people_loc)

    # 热力图绘制底图
    # img_square_app = cv2.applyColorMap(img_square_org, 1)
    img_square_app = cv2.imread('img_square_app.png')

    frame_count = 0
    #  随机走动
    while True:
        draw_map = img_square_signs.copy()
        draw_app_map = img_square_app.copy()

        draw_point_increase(draw_app_map)
        draw_app_map = apply_heatmap(draw_app_map, obs_people_loc)  # 按照obs路标绘制heatmap

        frame_count += 1
        if frame_count == 2:
            # 障碍人群
            for obs in obs_people_loc:
                """(1) 绘制障碍人群"""
                cv2.circle(draw_map, (obs[0], obs[1]), 7, (0, 0, 255), -1)  # （深红）

                if obs in safety_loc:
                    continue

                """(2) 更新障碍位置"""
                nex_loc = road_signs_neighbor[str(obs)]['neighbour'][random.randint(0, len(road_signs_neighbor[str(obs)]['neighbour']) - 1)]
                obs_people_loc_new.append(nex_loc)
        else:
            for obs in obs_people_loc:
                """(1) 绘制障碍人群"""
                cv2.circle(draw_map, (obs[0], obs[1]), 7, (0, 0, 255), -1)  # （深红）

                if obs in safety_loc:
                    continue

        # 客户人群
        for j, people in enumerate(people_loc):
            """(1) 绘制客户人群"""
            cv2.circle(draw_map, (people[0], people[1]), 7, [226, 43, 138], -1)  # （深紫）
            cv2.circle(draw_app_map, (people[0], people[1]), 7, [226, 43, 138], -1)  # （深紫）

        """（3）更新客户位置"""
        for people in people_loc:
            nex_loc = road_signs_neighbor[str(people)]['neighbour'][random.randint(0, len(road_signs_neighbor[str(people)]['neighbour']) - 1)]
            if (nex_loc in obs_people_loc) or (nex_loc in people_loc):  # or (best_path[1] in obs_people_loc_new)
                people_loc_new.append(people)
            else:
                people_loc_new.append(nex_loc)

        people_loc = copy.deepcopy(people_loc_new)
        people_loc_new.clear()
        if frame_count == 2:
            frame_count = 0
            obs_people_loc = copy.deepcopy(obs_people_loc_new)
            obs_people_loc_new.clear()

        # 地图更新
        cv2.imshow('draw_map', draw_map)
        cv2.imshow('img_app', draw_app_map)
        key = cv2.waitKey(1)

        # 暂停键
        if key == ord(' '):
            k = cv2.waitKey()
            if k == ord(' '):
                continue
        elif key == ord('q'):
            break

    # 开始导航
    frame_count = 0
    while reach_count != people_group_count:
        draw_map = img_square_signs.copy()
        draw_app_map = img_square_app.copy()

        # 障碍人群
        draw_point_increase(draw_app_map)
        draw_app_map = apply_heatmap(draw_app_map, obs_people_loc)

        frame_count += 1
        if frame_count == 2:
            for obs in obs_people_loc:
                """(1) 绘制障碍人群"""
                cv2.circle(draw_map, (obs[0], obs[1]), 7, (0, 0, 255), -1)  # （深红）

                if obs in safety_loc:
                    continue

                """(2) 更新障碍位置"""
                obs_neighbors = road_signs_neighbor_org[str(obs)]['neighbour']
                length_to_safety, min_length, next_loc = float('inf'), float('inf'), obs
                min_dis, min_ind = float('inf'), 0

                for i, neighbor in enumerate(obs_neighbors):
                    if neighbor in safety_loc:  # 移动策略一：向近在眼前的safety移动
                        next_loc = neighbor
                        break
                    min_safety = safety_loc[np.argmin([Graph.heuristic(neighbor, safety) for safety in safety_loc])]
                    length_to_safety = len(path_find(neighbor, min_safety))
                    if length_to_safety <= 7:
                        if length_to_safety < min_length:
                            min_length = length_to_safety
                            next_loc = neighbor

                    else:  # 移动策略二：向附近的其他人群靠拢
                        for o in obs_people_loc:
                            if o != obs:
                                dis = Graph.heuristic(neighbor, o)
                                if dis < min_ind:
                                    min_dis = dis
                                    min_ind = i
                        next_loc = obs_neighbors[min_ind]

                if (next_loc in obs_people_loc) or (next_loc in obs_people_loc_new) or (next_loc in obs_people_loc_new):
                    obs_people_loc_new.append(obs)
                else:
                    obs_people_loc_new.append(next_loc)
        else:
            for obs in obs_people_loc:
                """(1) 绘制障碍人群"""
                cv2.circle(draw_map, (obs[0], obs[1]), 7, (0, 0, 255), -1)  # （深红）

                if obs in safety_loc:
                    continue

        # 客户人群
        for j, people in enumerate(people_loc):
            """(1) 绘制客户人群"""
            cv2.circle(draw_map, (people[0], people[1]), 7, [226, 43, 138], -1)  # （深紫）
            cv2.circle(draw_app_map, (people[0], people[1]), 7, [226, 43, 138], -1)  # （深紫）

            if people in safety_loc:  # (若客户已到达safety)
                reach_count += 1

            else:  # (若客户还未到达safety)
                """(2) 找到用时最短路径"""
                road_signs_neighbor = update_road_signs_neighbor(obs_people_loc, road_signs_neighbor_org)  # 更新road_signs_neighbor
                dis, paths = [], []
                for safety in safety_loc:
                    path, time_length = path_find(people, safety), 0
                    for i in range(0, len(path) - 1):
                        index = road_signs_neighbor[str(path[i])]['neighbour'].index(path[i + 1])
                        time_length += road_signs_neighbor[str(path[i])]['time'][index]
                    dis.append(time_length)
                    paths.append(path)
                best_path = paths[dis.index(min(dis))]
                draw_path(draw_map, best_path)  # 绘制路径
                draw_path(draw_app_map, best_path)

                """（3）更新客户位置"""
                next_loc = best_path[1]
                if next_loc in people_loc:  # (next_loc in obs_people_loc) or (next_loc in obs_people_loc_new)
                    people_loc_new.append(people)
                else:
                    people_loc_new.append(next_loc)
        people_loc = copy.deepcopy(people_loc_new)
        people_loc_new.clear()
        if frame_count == 2:
            frame_count = 0
            obs_people_loc = copy.deepcopy(obs_people_loc_new)
            obs_people_loc_new.clear()

        # 地图更新
        cv2.imshow('draw_map', draw_map)
        cv2.imshow('img_app', draw_app_map)
        key = cv2.waitKey(1)

        # 暂停键
        if key == ord(' '):
            k = cv2.waitKey()
            if k == ord(' '):
                continue
        elif key == ord('q'):
            return


